#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI 脚本调度器主程序
"""

import sys
import os
from pathlib import Path
import subprocess
# 先设置sys.path，确保能找到public目录
current_dir = Path(__file__).parent  # py_scripts目录
if str(current_dir) not in sys.path:
    sys.path.insert(0, str(current_dir))

def ensure_package(pkg_name, import_name=None):
    """
    检查并自动安装指定包
    :param pkg_name: pip 包名
    :param import_name: import 时的模块名（如不同于pip包名）
    """
    import_name = import_name or pkg_name
    try:
        __import__(import_name)
        print(f"✅ {pkg_name} 已安装")
    except ImportError:
        print(f"📦 安装 {pkg_name}...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", pkg_name])
            print(f"✅ {pkg_name} 安装完成")
        except subprocess.CalledProcessError as e:
            print(f"❌ {pkg_name} 安装失败: {e}")
            return False
    return True

def install_required_packages():
    """安装必需的依赖包"""
    print("📦 检查并安装必需的依赖包...")
    
    # 依赖包列表
    REQUIRED_PACKAGES = [
        ("requests", None),
        ("huggingface_hub", None),
        ("packaging", None),
        ("setuptools", None),
        ("psutil", None),
    ]
    
    success = True
    for pkg, import_name in REQUIRED_PACKAGES:
        if not ensure_package(pkg, import_name):
            success = False
    
    if success:
        print("✅ 所有依赖包检查完成")
    else:
        print("⚠️ 部分依赖包安装失败，但程序将继续运行")
    
    return success

# 先安装依赖包
install_required_packages()


# 现在可以安全导入env_setup
from public.env_setup import setup_environment, ENV_CONFIG

# 使用统一的环境配置
ENV_CONFIG = setup_environment()

def init_environment():
    """初始化环境"""
    # 使用统一的环境配置，不需要手动设置sys.path
    print("✅ 环境初始化完成")
    print(f"📂 当前目录: {ENV_CONFIG['BASE_DIR']}")
    print(f"🐍 Python路径: {sys.executable}")
    

def load_plugins():
    """加载所有插件"""
    try:
        from public import PLUGINS
        if not PLUGINS:
            print("❌ 没有找到可用的插件")
            return False
        print("✅ 插件加载完成")
        return True
    except Exception as e:
        print(f"❌ 插件加载失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def show_menu():
    """显示插件菜单"""
    from public import PLUGINS

    print("\n" + "="*60)
    print("                ComfyUI 脚本调度器")
    print("="*60)

    menu_items = []
    for k, v in PLUGINS.items():
        desc = v.__doc__ or ""
        menu_items.append({'id': k, 'name': k, 'description': desc})

    if not menu_items:
        print("❌ 没有找到可用的插件")
        return []

    for i, item in enumerate(menu_items, 1):
        print(f"{i:2d}. 📦 {item['name']}")
        if item['description']:
            print(f"    {item['description']}")

    print(f"{len(menu_items)+1:2d}. 🚪 退出")
    print("="*60)

    return menu_items

def main():
    """主函数"""
    print("🚀 启动 ComfyUI 脚本调度器...")
    
    # 初始化环境
    init_environment()
    
    # 加载插件
    if not load_plugins():
        print("❌ 插件加载失败，程序退出")
        input("按任意键退出...")
        return
    
    # 主循环
    while True:
        menu_items = show_menu()
        
        if not menu_items:
            print("按任意键退出...")
            input()
            break
        
        try:
            choice = input(f"\n请选择功能 (1-{len(menu_items)+1}): ").strip()
            
            if choice == str(len(menu_items)+1):
                print("👋 再见！")
                break
                
            if 1 <= int(choice) <= len(menu_items):
                selected_item = menu_items[int(choice)-1]
                plugin_id = selected_item['id']
                
                print(f"\n🚀 执行: {selected_item['name']}")
                print("-" * 40)
                
                # 运行插件
                from public import PLUGINS
                plugin_cls = PLUGINS.get(plugin_id)
                if plugin_cls:
                    plugin = plugin_cls()
                    if hasattr(plugin, 'start'):
                        plugin.start()
                    elif hasattr(plugin, 'main'):
                        plugin.main()
                    else:
                        print("插件没有可执行入口")
                        success = False
                    success = True
                else:
                    print("未找到插件")
                    success = False
                
                if success:
                    print("✅ 执行完成")
                else:
                    print("❌ 执行失败")
                    
                print("-" * 40)
                input("按回车键返回主菜单...")
            else:
                print("❌ 无效选项")
                
        except KeyboardInterrupt:
            print("\n👋 用户中断，再见！")
            break
        except ValueError:
            print("❌ 请输入有效的数字")
        except Exception as e:
            print(f"❌ 错误: {e}")

if __name__ == "__main__":
    main() 